import shutil
import traceback

import nrrd
import cv2
import scipy.ndimage as ndi
import csv

import tqdm

from monai.data.image_reader import pydicom

from models.experimental import attempt_load  # 加载 YOLOv7 的架构
from utils.general import non_max_suppression, scale_coords
import torch
from PIL import Image
import numpy as np
from nnunet.inference.predict import predict_from_folder
import nibabel as nib
import os
import SimpleITK as sitk
from skimage.measure import label, regionprops
from glob import glob
from pypinyin import lazy_pinyin
import zipfile
import re
from totalsegmentator.python_api import totalsegmentator


def unzip(input_folder, output_folder):
    zip_files = [f for f in os.listdir(input_folder) if f.endswith('zip')]
    for zip_file in zip_files:
        zip_path = os.path.join(input_folder, zip_file)
        output = os.path.join(output_folder, zip_file)
        makedir_folder(output)
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(output)


def dicom_series_to_nrrd(input_folder, output, csv_path):
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        write = csv.writer(file)
        write.writerow(['nrrd_name'])
        for root, dirs, files in os.walk(input_folder):
            try:
                for file in files:
                    path = os.path.join(root, file)
                    file_path = path.replace('/' + os.path.basename(path), "")
                    ct_name = os.path.basename(file_path)
                    ID = os.path.dirname(os.path.dirname(os.path.dirname(path)))
                    ID_number = os.path.basename(ID)

                    for image in os.listdir(file_path):
                        image_path = os.path.join(file_path, image)
                        ds = pydicom.dcmread(image_path)
                        patient_ID = ds.PatientID
                        patient_name = ds.PatientName
                        patient_birth = ds.PatientBirthDate
                        continue
                    patient_name = ''.join(lazy_pinyin(patient_name))
                    patient_name = re.sub(r'[^a-zA-Z]', '', patient_name)

                    if '/' in patient_ID or '/' in patient_name or '/' in patient_birth:
                        continue
                    name = str(ID_number) + str(patient_name).replace(' ', '') + '_' + str(patient_birth) + '_' + str(
                        patient_ID) + '_' + ct_name + '.nrrd'
                    output_nrrd = os.path.join(output, name)
                    if os.path.exists(output_nrrd):
                        break
                    print(output_nrrd)
                    # 使用SimpleITK中ImageSeriesReader类的GetGDCMSeriesFileNames方法获取指定文件夹中DICOM系列的文件名列表。
                    dicom_names = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(file_path)
                    # 实例化了一个ImageSeriesReader对象，它用于按顺序读取DICOM文件
                    reader = sitk.ImageSeriesReader()
                    # 设置了reader对象的文件名列表，以便读取DICOM信息
                    reader.SetFileNames(dicom_names)
                    # 调用Execute方法读取DICOM系列，将结果保存在dicom_series中
                    dicom_series = reader.Execute()
                    # 使用SimpleITK的WriteImage方法将dicom_series写入到指定的NRRD文件中-

                    sitk.WriteImage(dicom_series, output_nrrd)
            except Exception as e:
                write.writerow([ID_number])


def run_totalseg(input_folder, output, save_nii, heart_nii):
    folder = os.path.normpath(input_folder)
    for file_folder in os.listdir(folder):
        file_path = os.path.join(folder, file_folder)
        nrrd_files = glob(file_path + '/*.nrrd')
        for nrrd_file in nrrd_files:
            name = os.path.basename(nrrd_file.replace('nrrd', 'nii.gz'))
            heart_path = os.path.join(heart_nii, name)
            # problem_file=os.path.join(problem_folder,nrrd_file)
            # if not os.path.exists(heart_path):
            if os.path.exists(heart_path):
                continue
            img = sitk.ReadImage(nrrd_file)
            data = os.path.basename(nrrd_file)
            data_name = os.path.splitext(data)[0]
            if ' ' in data_name:
                data_name = data_name.replace(' ', '')
            data_name_nii = data_name + '.nii.gz'
            save_path = os.path.join(save_nii, data_name_nii)
            sitk.WriteImage(img, save_path)

            output_folder = os.path.join(output, data_name)
            if os.path.exists(output_folder):
                continue
            os.makedirs(output_folder, exist_ok=True)
            try:
                command_line = "/home/imed/Development/Tools/Slicer-5.6.2-linux-amd64/bin/PythonSlicer /home/imed/Development/Tools/Slicer-5.6.2-linux-amd64/lib/Python/bin/TotalSegmentator -i {} -o {} --task total"
                new_command_line = command_line.format(save_path, output_folder)
                # print(os.system(new_command_line))
                totalsegmentator(save_path, output_folder, task="total", fast=False, verbose=True)
            except Exception as e:
                print(traceback.format_exc())


def niigz_to_slice_arota(input_folder, output_folder):
    for folder in os.listdir(input_folder):
        input_path_folder = os.path.join(input_folder, folder)
        output_path = os.path.join(output_folder, folder)
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        for file in os.listdir(input_path_folder):
            if 'aorta' in file:
                input_path = os.path.join(input_path_folder, file)
                img = nib.load(input_path)
                data = img.get_fdata()
                shape = data.shape
                for i in range(shape[2]):
                    slice_data = data[:, :, i]
                    min_val = np.min(slice_data)
                    max_val = np.max(slice_data)
                    if max_val - min_val == 0:
                        slice_data = np.zeros_like(slice_data, dtype=np.uint8)
                    else:
                        slice_data = np.uint8(255 * (slice_data - min_val) / (max_val - min_val))

                    slice_img = Image.fromarray(slice_data)
                    slice_img.save(os.path.join(output_path, f'slice_{i:04d}.png'))


def expand_white_regions(input_folder, dilations_size=3):
    for root, dir, files in os.walk(input_folder):
        for file in files:
            image_path = os.path.join(root, file)
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            _, binary_image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
            if np.all(binary_image == 0):
                continue
            print(image_path)
            num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary_image, connectivity=8)
            result = np.copy(binary_image)
            kernel = np.ones((dilations_size, dilations_size), np.uint8)
            for label in range(1, num_labels):
                component_mask = (labels == label).astype(np.uint8) * 255
                inner_component = cv2.erode(component_mask, kernel, iterations=1)
                result = cv2.bitwise_or(result, inner_component)
            cv2.imwrite(image_path, result)


def apply_condition_to_image_arota(folder1, folder2, save_path):
    for file in os.listdir(folder1):
        folder1_file = os.path.join(folder1, file)
        # name=file.replace('CCTA_NC',' CCTA_NC')
        name = file + '.nrrd'
        folder2_file = os.path.join(folder2, name)
        save_folder = os.path.join(save_path, name)
        makedir_folder(save_folder)
        for image in os.listdir(folder1_file):
            image1_path = os.path.join(folder1_file, image)
            image2_path = os.path.join(folder2_file, image)
            makedir_folder(image2_path)
            save = os.path.join(save_folder, image)
            img1 = cv2.imread(image1_path, cv2.IMREAD_GRAYSCALE)
            img2 = cv2.imread(image2_path, cv2.IMREAD_GRAYSCALE)
            _, binary1 = cv2.threshold(img1, 127, 255, cv2.THRESH_BINARY)
            _, binary2 = cv2.threshold(img2, 127, 255, cv2.THRESH_BINARY)
            num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary2, connectivity=8)
            result = binary2.copy()
            img3 = binary2.copy()
            for label in range(1, num_labels):
                component_mask = (labels == label).astype(np.uint8) * 255
                overlap = cv2.bitwise_and(component_mask, binary1)
                if np.any(overlap):
                    result[labels == label] = 0
            cv2.imwrite(save, result)


def aorta_cac_filter(folder1, folder2, save_path):
    for file in os.listdir(folder1):
        folder1_file = os.path.join(folder1, file)
        # name=file.replace('CCTA_NC',' CCTA_NC')
        name = file + '.nrrd'
        folder2_file = os.path.join(folder2, name)
        save_folder = os.path.join(save_path, name)
        makedir_folder(save_folder)

        for image in os.listdir(folder1_file):
            image1_path = os.path.join(folder1_file, image)
            image2_path = os.path.join(folder2_file, image)
            makedir_folder(image2_path)
            save = os.path.join(save_folder, image)
            img1 = cv2.imread(image1_path, cv2.IMREAD_GRAYSCALE)
            img2 = cv2.imread(image2_path, cv2.IMREAD_GRAYSCALE)
            _, binary1 = cv2.threshold(img1, 127, 255, cv2.THRESH_BINARY)
            _, binary2 = cv2.threshold(img2, 127, 255, cv2.THRESH_BINARY)
            img3 = np.zeros_like(img2)
            labeled_img2 = label(img2)
            regions2 = regionprops(labeled_img2)

            for region in regions2:
                minr, minc, maxr, maxc = region.bbox
                if np.any(img1[minr:maxr, minc:maxc] == 255):
                    img3[minr:maxr, minc:maxc] = img2[minr:maxr, minc:maxc]

            cv2.imwrite(save, img3)


def convert_nrrd_to_nii(input_folder, output_folder):
    for root, dir, files in os.walk(input_folder):
        for file in files:
            # if 'CCTA_NC.nrrd' in file:
            input_path = os.path.join(root, file)
            name = file.replace('.nrrd', '_0000.nii.gz')
            output_path = os.path.join(output_folder, name)
            if os.path.exists(output_path):
                continue
            image = sitk.ReadImage(input_path)
            sitk.WriteImage(image, output_path)


def setup_nnunet_env():
    os.environ["nnUNet_raw_data_base"] = "/mnt/2093089d-5770-4686-b3f4-686eb07a909e/heart_1204_v1/heart/124"
    os.environ["nnUNet_preprocessed"] = "/mnt/2093089d-5770-4686-b3f4-686eb07a909e/heart_1204_v1/heart/processed"
    os.environ["RESULTS_FOLDER"] = "/mnt/2093089d-5770-4686-b3f4-686eb07a909e/heart_1204_v1/heart"
    for var in ["nnUNet_raw_data_base", "nnUNet_preprocessed", "RESULTS_FOLDER"]:
        if var not in os.environ:
            raise EnvironmentError(f"{var} not in environ")


def run_nnunet_inference(input_path, nnunet_folder, output_path, csv_path, folds=None):
    if folds is None:
        folds = [2]
    model_folder = os.environ["RESULTS_FOLDER"]
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file1:
        write = csv.writer(file1)
        write.writerow(['file_name', 'total_HU', 'area'])
        for file in os.listdir(input_path):
            heart_outcome_path = os.path.join(output_path, file).replace('_0000.nii.gz', '.nii.gz')
            if os.path.exists(heart_outcome_path):
                print(heart_outcome_path)
                continue
            file_path = os.path.join(input_path, file)
            nnunet_path = os.path.join(nnunet_folder, file)
            shutil.copy2(file_path, nnunet_path)
            try:
                predict_from_folder(
                    lowres_segmentations=None,
                    part_id=0,
                    num_parts=1,
                    tta=False,
                    input_folder=nnunet_folder,
                    output_folder=output_path,
                    model=model_folder,
                    folds=folds,
                    save_npz=False,
                    num_threads_preprocessing=1,
                    num_threads_nifti_save=1,
                    overwrite_existing=True,
                    segmentation_export_kwargs=None
                )

            except Exception as e:
                write.writerow(file)
                print(file_path)
            os.remove(nnunet_path)


def niigz_to_slice(input_folder, output_folder):
    for file in os.listdir(input_folder):
        if 'plans' not in file:
            input_path = os.path.join(input_folder, file)
            output_path = os.path.join(output_folder, file)
            if not os.path.exists(output_path):
                os.makedirs(output_path)
            img = nib.load(input_path)
            data = img.get_fdata()
            shape = data.shape
            for i in range(shape[2]):
                slice_data = data[:, :, i]
                slice_data = np.uint8(
                    255 * (slice_data - np.min(slice_data)) / (np.max(slice_data) - np.min(slice_data)))
                slice_img = Image.fromarray(slice_data)
                slice_img.save(os.path.join(output_path, f'slice_{i:04d}.png'))


def get_heart_roi(ori_path, output):
    # 读取原始和标签图像
    data, data_header = nrrd.read(ori_path)

    for i in range(data.shape[2]):  # 假设数据是 (height, width, depth)
        slice_data = data[:, :, i]
        slice_image = Image.fromarray(slice_data)
        slice_image = slice_image.convert('L')  # 如果数据是灰度图
        slice_image.save(os.path.join(output, f'slice_{i:04d}.png'))


def apply_condition_to_image(folder1, folder2, save_path):
    for file in os.listdir(folder1):
        folder1_file = os.path.join(folder1, file)
        name = file.replace('nii.gz', 'nrrd')
        folder2_file = os.path.join(folder2, name)
        save_folder = os.path.join(save_path, name)
        makedir_folder(save_folder)
        for image in os.listdir(folder1_file):
            image1_path = os.path.join(folder1_file, image)
            image2_path = os.path.join(folder2_file, image)
            save = os.path.join(save_folder, image)
            img1 = Image.open(image1_path).convert('L')
            img2 = Image.open(image2_path).convert('L')
            img1_array = np.array(img1)
            img2_array = np.array(img2)
            output_array = np.zeros_like(img2_array)
            mask_white_pixel = (img1_array > 0)
            output_array[mask_white_pixel] = img2_array[mask_white_pixel]
            output_image = Image.fromarray(output_array)
            output_image.save(save)


def get_lable_slice(label_nrrd, label_slice, num):
    data, header = nrrd.read(label_nrrd)
    data = data[num, :, :, :]
    # 确保输出目录存在
    output_dir = label_slice
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 保存每个切片
    for i in range(data.shape[2]):  # 假设数据是 (height, width, depth)
        slice_data = data[:, :, i]
        slice_image = Image.fromarray(slice_data)
        slice_image = slice_image.convert('L')  # 如果数据是灰度图
        slice_image.save(os.path.join(output_dir, f'slice_{i:04d}.png'))


def process_image(image_path):
    # Open the image
    image = Image.open(image_path).convert('L')  # Convert to grayscale
    # Convert image to numpy array
    image_array = np.array(image)
    # Create a mask where pixels with value 1 are set to 255, and others to 0
    processed_array = np.where(image_array == 1, 255, 0).astype(np.uint8)
    # Convert the numpy array back to an image
    processed_image = Image.fromarray(processed_array)
    return processed_image


def get_heart_label(label_slice):
    # Iterate over all files in the directory
    for filename in os.listdir(label_slice):
        if filename.endswith('.png') or filename.endswith('.jpg') or filename.endswith('.jpeg'):
            # Construct the full file path
            file_path = os.path.join(label_slice, filename)
            # Process the image
            processed_image = process_image(file_path)
            # Save the processed image
            processed_image.save(os.path.join(label_slice, f'{filename}'))


# Function to binarize a single image
def binarize_image(image_path, threshold=200):
    # Open the image
    image = Image.open(image_path).convert('L')  # Convert to grayscale
    # Convert image to numpy array
    image_array = np.array(image)
    # Binarize the image
    binary_array = np.where(image_array > threshold, 255, 0).astype(np.uint8)
    # Convert the numpy array back to an image
    binary_image = Image.fromarray(binary_array)
    return binary_image


def get_cac(input_path, heart_slice, threshold=140):
    # Iterate over all files in the directory
    for filename in os.listdir(input_path):
        if filename.endswith('.png') or filename.endswith('.jpg') or filename.endswith('.jpeg'):
            # Construct the full file path
            file_path = os.path.join(input_path, filename)
            # Binarize the image
            binary_image = binarize_image(file_path, threshold)
            # Save the binarized image
            binary_image.save(os.path.join(heart_slice, filename))


def get_slice_1(folder, heart_niigz):
    for root, dir, files in os.walk(folder):
        for file in files:
            name = file.replace('nrrd', 'nii.gz')
            heart_path = os.path.join(heart_niigz, name)
            if not os.path.exists(heart_path):
                continue
            heart_image_folder = os.path.join(heart_image, file)
            new_label_folder = os.path.join(new_label, file)
            makedir_folder(new_label_folder)
            get_cac(heart_image_folder, new_label_folder)


def load_model(weights_path='best.pt', device='cuda'):
    # 加载模型架构及权重
    model = attempt_load(weights_path, map_location=device)
    model.eval()  # 设置为评估模式
    return model


# 2. 加载并预处理图片
def load_image(image_path, img_size=640):
    image = cv2.imread(image_path)  # 使用 OpenCV 读取图片
    original_image = image.copy()  # 保留原图用于后续绘制结果
    image = cv2.resize(image, (img_size, img_size))  # 将图片调整到 640x640
    image = image[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB and HWC to CHW
    image = np.ascontiguousarray(image)  # 使数组在内存中连续
    image = torch.from_numpy(image).float() / 255.0  # 转换为张量并归一化
    image = image.unsqueeze(0)  # 增加批次维度
    return image, original_image


# 3. 对图片进行推理
def infer_image(model, image, device='cuda'):
    image = image.to(device)  # 将图片移动到 GPU 或 CPU
    with torch.no_grad():  # 禁用梯度计算，以提高推理速度
        outputs = model(image)[0]  # 使用模型进行推理，获取第一个输出（检测框）
    return outputs


# 4. 进行非极大值抑制 (NMS) 处理
def process_detections(outputs, conf_threshold=0.25, iou_threshold=0.45):
    # 使用非极大值抑制来消除重叠框
    results = non_max_suppression(outputs, conf_threshold, iou_threshold)
    return results


# 5. 在图片上绘制检测结果
def draw_detections(original_image, image, results, img_size=640):
    for result in results:
        if result is not None:
            result[:, :4] = scale_coords(image.shape[2:], result[:, :4], original_image.shape).round()
            for detection in result:
                x1, y1, x2, y2, conf, cls = detection[:6].cpu().numpy()
                # 绘制检测框
                cv2.rectangle(original_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 1)
                # 显示类别和置信度
                cv2.putText(original_image, f' Conf: {conf:.2f}',
                            (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)

    return original_image


def yolo_detect(folders, output, weight):
    weights_path = weight  # 替换为你的模型权重路径
    img_size = 640  # 图片输入尺寸
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = load_model(weights_path, device)

    for folder in os.listdir(folders):
        folder_path = os.path.join(folders, folder)
        output_path = os.path.join(output, folder)
        makedir_folder(output_path)
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        for file in os.listdir(folder_path):
            image_path = os.path.join(folder_path, file)
            output_image = os.path.join(output_path, file)

            # 加载图片并进行预处理
            image, original_image = load_image(image_path, img_size)

            # 推理
            outputs = infer_image(model, image, device)

            # 后处理：非极大值抑制
            results = process_detections(outputs)

            # 在图片上绘制检测框
            result_image = draw_detections(original_image, image, results)

            # 保存结果图片
            cv2.imwrite(output_image, result_image)


def find_green_boxes(image):
    # 假设绿色框的颜色在 BGR 空间中的范围
    lower_green = np.array([0, 200, 0])  # BGR 颜色范围
    upper_green = np.array([100, 255, 100])

    # 创建掩模
    mask = cv2.inRange(image, lower_green, upper_green)

    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    boxes = []
    for contour in contours:
        # 获取边界框坐标
        x, y, w, h = cv2.boundingRect(contour)
        boxes.append((x, y, x + w, y + h))  # (x1, y1, x2, y2)

    return boxes


def filter_white_pixels(base_image, green_boxes):
    # 创建一个和 base_image 相同的全黑图像
    result_image = np.zeros_like(base_image)

    for (x1, y1, x2, y2) in green_boxes:
        # 提取绿色框内的区域
        region = base_image[y1:y2, x1:x2]

        # 保留白色像素
        white_pixels = (region == 255)  # 假设白色像素值为 255
        result_image[y1:y2, x1:x2][white_pixels] = region[white_pixels]

    return result_image


def filter_liantongyu(folder, output):
    for file in os.listdir(folder):
        image_folder = os.path.join(folder, file)
        out_folder = os.path.join(output, file)
        if not os.path.exists(out_folder):
            os.makedirs(out_folder)
        for image in os.listdir(image_folder):
            image_path = os.path.join(image_folder, image)
            output_path = os.path.join(out_folder, image)

            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 读取为灰度图像
            # 确保图像是二值化的（白色像素为255，黑色像素为0）
            _, binary_image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
            # 查找连通域
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_image, connectivity=8)
            num_components = num_labels - 1
            if num_components > 6:
                threshold = 6
            else:
                threshold = 4
            # 创建一个空的图像用于存储结果
            processed_image = np.zeros_like(image)
            # 遍历所有连通域
            for i in range(1, num_labels):  # 从1开始，0是背景
                area = stats[i, cv2.CC_STAT_AREA]
                if area > threshold:
                    processed_image[labels == i] = 255
            cv2.imwrite(output_path, processed_image)
            print(output_path)


def filter_yolo(new_label_folder, yolo_label_folder, end_label_folder):
    for folder in os.listdir(new_label_folder):
        new_label_path = os.path.join(new_label_folder, folder)
        yolo_label_path = os.path.join(yolo_label_folder, folder)
        end_label_path = os.path.join(end_label_folder, folder)
        if not os.path.exists(end_label_path):
            os.makedirs(end_label_path)
        for image in os.listdir(new_label_path):
            new_label_image = os.path.join(new_label_path, image)
            yolo_label_image = os.path.join(yolo_label_path, image)
            end_label_image = os.path.join(end_label_path, image)
            image_with_boxes = cv2.imread(yolo_label_image)  # 包含绿色框的图像
            base_image = cv2.imread(new_label_image)  # 另一张图像
            green_boxes = find_green_boxes(image_with_boxes)

            # 在另一张图上保留白色像素
            filtered_image = filter_white_pixels(base_image, green_boxes)

            # 保存结果图像
            cv2.imwrite(end_label_image, filtered_image)


def find_matching_files(folderA, folderB):
    for root, dirs, files in os.walk(folderB):
        for file in files:
            if file == folderA:
                filename = os.path.join(root, file)
                return os.path.join(root, file)
    return None


def reconstruct_3d_image(image_dir, output_nrrd_path, original_nrrd_path):
    _, header = nrrd.read(original_nrrd_path)
    image_files = sorted([os.path.join(image_dir, f) for f in os.listdir(image_dir)])
    slices = []
    for image_file in image_files:
        image_path = image_file
        image = Image.open(image_path).convert("L")
        image_array = np.array(image)
        slices.append(image_array)
    print(output_nrrd_path)
    volume = np.stack(slices, axis=-1)
    print(volume.shape)
    nrrd.write(output_nrrd_path, volume, header=header)


def to3D(heart_slice_dir, segment_path, output_nrrd_file, heart_niigz):
    for root, dirs, files in os.walk(segment_path):
        for file in files:
            try:
                name = file.replace('nrrd', 'nii.gz')
                heart_path = os.path.join(heart_niigz, name)
                if not os.path.exists(heart_path):
                    continue

                old_dir_path = os.path.join(heart_slice_dir, file)

                segment_file = os.path.join(root, file)
                output_path = os.path.join(output_nrrd_file, file)
                reconstruct_3d_image(old_dir_path, output_path, segment_file)
            except Exception as e:
                print(file)


def calculate_HU(folder_path, image_path, csv_path, heart_niigz):
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        write = csv.writer(file)
        write.writerow(['file_name', 'total_HU', 'area'])
        for root, dir, files in os.walk(image_path):
            for file_name in files:
                try:

                    name = file_name.replace('nrrd', 'nii.gz')
                    heart_path = os.path.join(heart_niigz, name)
                    if not os.path.exists(heart_path):
                        continue

                    file_name1 = file_name
                    file_path = os.path.join(folder_path, file_name1)
                    image = sitk.ReadImage(file_path)
                    image_array = sitk.GetArrayFromImage(image)
                    old_image_path = os.path.join(root, file_name)
                    print(old_image_path)
                    image_old = sitk.ReadImage(old_image_path)
                    image_array_old = sitk.GetArrayFromImage(image_old)
                    total = 0.0
                    area = 0.0
                    for z in range(image_array.shape[0]):
                        # 计算面积
                        label_images, num_features = ndi.label(image_array[z, :, :])
                        slicer_total = 0
                        cac_area = 0
                        for label in range(1, num_features + 1):
                            mask = np.zeros_like(image_array[z, :, :])
                            mask[label_images == label] = 255
                            slice_data = mask
                            white_pixel_count = np.sum(slice_data == 255)
                            if white_pixel_count < 4:
                                continue
                            data, header = nrrd.read(old_image_path)
                            pixel_spacing = header.get('space directions', None)
                            pixel_area_slice = np.abs(pixel_spacing[0, 0] * pixel_spacing[1, 1])
                            white_pixel_area = white_pixel_count * pixel_area_slice
                            # 计算weight
                            ct_data = image_array_old[z, :, :]
                            indices_255 = np.argwhere(slice_data == 255)
                            if len(indices_255) > 0:
                                hu_value = [ct_data[i, j] for i, j in indices_255]
                                max_hu_index = indices_255[np.argmax(hu_value)]
                            hu_value = np.abs(image_array_old[z, max_hu_index[0], max_hu_index[1]])
                            if 130 <= hu_value < 200:
                                weight = 1
                            elif 200 <= hu_value < 300:
                                weight = 2
                            elif 300 <= hu_value < 400:
                                weight = 3
                            elif hu_value >= 400:
                                weight = 4
                            slicer_total += weight * white_pixel_area
                            cac_area += white_pixel_area
                        total += slicer_total
                        area += cac_area
                    print(total, area)
                    write.writerow([file_name, total, area])
                except Exception as e:
                    print(file_name)


def makedir_folder(folder):
    if not os.path.exists(folder):
        os.makedirs(folder)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

"""""""""""""""""""""""""""""""""""""""  基础文件夹定义  """""""""""""""""""""""""""""""""""""""""""""""
dataset_base_dir = "/home/imed/Development/Tmps/Dataset_8000_Test"

zip_folder = '/mnt/2093089d-5770-4686-b3f4-686eb07a909e/heart_1204_v1/heart/zip'
input_folder = os.path.join(dataset_base_dir, "213")

folder = os.path.join(dataset_base_dir, "2")
folder1 = os.path.join(folder, "ori_data")
makedir_folder(folder1)

nnunet_folder = os.path.join(folder, "nnunet_folder")
makedir_folder(nnunet_folder)

ori_path = os.path.join(dataset_base_dir, 'v11_CCTA_NC/ori_heart_slice')
makedir_folder(ori_path)

heart_filter = os.path.join(dataset_base_dir, 'v11_CCTA_NC/heart_filter')
heart_image = os.path.join(dataset_base_dir, 'v11_CCTA_NC/heart_image')
label_folder = os.path.join(dataset_base_dir, 'v11_CCTA_NC/old_label')
new_label = os.path.join(dataset_base_dir, 'v11_CCTA_NC/newlabel')
yolo_detect_output = os.path.join(dataset_base_dir, 'v11_CCTA_NC/yolo_detect')
weight_cac = 'best.pt'
new2label = os.path.join(dataset_base_dir, 'v11_CCTA_NC/new2label')
new3label = os.path.join(dataset_base_dir, 'v11_CCTA_NC/new3label')
processed_image = os.path.join(dataset_base_dir, 'v11_CCTA_NC/new_segment')
processed_aorta_cac = os.path.join(dataset_base_dir, 'v11_CCTA_NC/aorta_cac_segment')
makedir_folder(processed_image)
makedir_folder(processed_aorta_cac)

old_heart_label = processed_aorta_cac = os.path.join(dataset_base_dir, 'v11_CCTA_NC/old_heart_label')
xinzang_nrrd = os.path.join(dataset_base_dir, 'v11_CCTA_NC/xinzang_nrrd')
makedir_folder(xinzang_nrrd)

niigz_foldr = os.path.join(dataset_base_dir, 'v11_CCTA_NC/niigz')
makedir_folder(niigz_foldr)

heart_niigz = os.path.join(dataset_base_dir, 'v11_CCTA_NC/heart_niigz')
makedir_folder(heart_niigz)

heart_niigz_slice = heart_niigz = os.path.join(dataset_base_dir, 'v11_CCTA_NC/heart_niigz_slice')
all_folder_path = os.path.join(dataset_base_dir, 'v11_CCTA_NC')
aorta = os.path.join(dataset_base_dir, 'v11_CCTA_NC/aorta')
makedir_folder(aorta)
save_total_nii = os.path.join(dataset_base_dir, 'v11_CCTA_NC/total_nii')
makedir_folder(save_total_nii)
aorta_slice = os.path.join(dataset_base_dir, 'v11_CCTA_NC/aorta_slice')

aorta_filter = os.path.join(dataset_base_dir, 'v11_CCTA_NC/aorta_filter')
aorta_cac = os.path.join(dataset_base_dir, 'v11_CCTA_NC/aorta_cac')
save_end = os.path.join(dataset_base_dir, 'results')
makedir_folder(save_end)

heart_cac_csv = os.path.join(save_end, 'heart_cac.csv')
aorta_csv = os.path.join(save_end, 'aorta_cac.csv')
probelm_csv = os.path.join(save_end, 'probelm_csv.csv')
problem_dicom_csv = os.path.join(save_end, 'problem_dicom_csv')

# 解压原始文件
# unzip(zip_folder,input_folder)
# dicom文件合成nrrd文件
# dicom_series_to_nrrd(input_folder, folder1, problem_dicom_csv)
# #nrrd原始文件转换为nii.gz文件
# convert_nrrd_to_nii(folder,niigz_foldr)
# 设置心脏分割模型环境
# setup_nnunet_env()
# 从原始图像分割心脏
# run_nnunet_inference(niigz_foldr,nnunet_folder,heart_niigz,probelm_csv)
# 将分割出来的心脏进行切片
# niigz_to_slice(heart_niigz,heart_niigz_slice)
# 将原始图像进行切片
# for root,dir,files in os.walk(folder):
#    for file in files:
#        name=file.replace('nrrd','nii.gz')
#        heart_path=os.path.join(heart_niigz,name)
#        if not os.path.exists(heart_path):
#            continue
#
#        nrrd_path=os.path.join(root,file)
#        ori_heart_slice=os.path.join(ori_path,file)
#        makedir_folder(ori_heart_slice)
#        get_heart_roi(nrrd_path,ori_heart_slice)
#
# for file in os.listdir(xinzang_nrrd):
#    file_path=os.path.join(xinzang_nrrd,file)
#    output_apth=os.path.join(heart_filter,file)
#    makedir_folder(output_apth)
#    get_heart_roi(file_path,output_apth)
# 从原始图像的切片中提取心脏区域
# apply_condition_to_image(heart_niigz_slice,ori_path,heart_image)


# 分割主动脉
run_totalseg(folder,aorta,save_total_nii,heart_niigz)
#将主动脉的分割结果进行切片
niigz_to_slice_arota(aorta,aorta_slice)
#expand_white_regions(aorta_slice)
#将原始图像的切片结果进行二值化，提取可能的钙化部分
get_slice_1(folder,heart_niigz)
#yolo在原始切片图像检测钙化物
yolo_detect(heart_image, yolo_detect_output, weight_cac)
#过滤像素个数过小的钙化物
filter_liantongyu(new_label,new2label)
#过滤yolo检测的非钙化物
filter_yolo(new2label,yolo_detect_output,new3label)
#将主动脉上的钙化物分离形成另外的图像
apply_condition_to_image_arota(aorta_slice,new3label,aorta_filter)
#过滤主动脉上的钙化物
aorta_cac_filter(aorta_slice,new3label,aorta_cac)
#合成冠脉钙化标签
to3D(aorta_filter, folder, processed_image,heart_niigz)
#合成主动脉钙化标签
to3D(aorta_cac,folder,processed_aorta_cac,heart_niigz)
#计算冠脉钙化积分
calculate_HU(processed_image, folder,heart_cac_csv,heart_niigz)
#计算主动脉钙化积分
calculate_HU(processed_aorta_cac,folder,aorta_csv,heart_niigz)


#将所有分割结果总结到一个文件夹
for file in os.listdir(folder1):
    try:
        name = file.replace('nrrd', 'nii.gz')
        heart_path = os.path.join(heart_niigz, name)
        if not os.path.exists(heart_path):
            continue
        ori_nrrd=os.path.join(folder1,file)
        name_aorta_cac=file.replace('.nrrd','aorta_cac.nrrd')
        aorta_cac=os.path.join(processed_aorta_cac,file)
        name_heart_cac=file.replace('.nrrd','heart_cac.nrrd')
        new_cac_seg=os.path.join(processed_image,file)
        name1=file.replace('nrrd','nii.gz')
        heart_seg=os.path.join(heart_niigz,name1)
        name2=file.replace('.nrrd','')
        path=os.path.join(aorta,name2)
        aorta_seg=os.path.join(path,'aorta.nii.gz')
        save_path=os.path.join(save_end,file)
        makedir_folder(save_path)
        shutil.copy2(ori_nrrd,os.path.join(save_path,file))
        shutil.copy2(aorta_cac, os.path.join(save_path, name_aorta_cac))
        shutil.copy2(new_cac_seg, os.path.join(save_path, name_heart_cac))
        shutil.copy2(heart_seg, os.path.join(save_path, name1))
        shutil.copy2(aorta_seg, os.path.join(save_path,'aorta.nii.gz' ))
    except Exception as e:
        print(file)
